# -*- coding: utf-8 -*-

## @package dmiParser
# @author: Ludovic Barriere
# @date: 07/02/2013
# @summary: This module is a toolbox of parsing functions 
# of input/output file in the \a DMI format 


import scipy.sparse as sp
from numpy import float
from numpy import argsort
import re


PCH_PATTERN_Key = re.compile(r"(........)(.*)")

PCH_HEAD = {"KEY":"DMI","NAME":None,"0":"0","FORM":None,"TIN":None,"TOUT":None,"M":None,"N":None}
PCH_HCOLUMN = {"KEY":"DMI","NAME":None,"J":None,"I":None,"A(I,J)":None,"align":"<"}
PCH_HCOLUMN_LONG = {"KEY":"DMI*","NAME":None,"J":None,"I":None,"A(I,J)":None,"align":"<"}
PCH_COLUMN = {"KEY":"        ","I":None,"A(I,J)":None,"align":">"}
PCH_COLUMN_LONG = {"KEY":"*       ","I":None,"A(I,J)":None,"align":">"}

HENTRY = "{KEY!s:<8}{NAME!s:<8}       0{FORM!s:>8}{TIN!s:>8}{TOUT!s:>8}        {M!s:>8}{N!s:>8}\n"

ENTRY = "{KEY!s:<8}{0[0]:<8}{0[1]:>8}{0[2]:>8}{0[3]:>8}{0[4]:>8}{0[5]:>8}{0[6]:>8}{0[7]:>8}\n"
ENTRY_LONG = "{KEY!s:<8}{0[0]:{align}16}{0[1]:>16}{0[2]:>16}{0[3]:>16}\n"

FLOAT = "float64"

## readPunch function: DMI parser
#  Read operators in a punch file in the DMI format
#  @param pch_file str(): name of punch file
#  @param output: sparse matrix class
#  @return operators dict(): dictionary of extracted operators

def readPunch(pch_file,output=sp.lil_matrix):
    """ readPunch documentation:
        
    >>> filename = "../test/test.pch"
    
    >>> OP = readPunch(filename,output=sp.csr_matrix)
    
    >>> OP['KGG1'].todense()
    matrix([[ 1.,  2.,  3.],
            [ 2.,  1.,  2.],
            [ 3.,  2.,  1.]], dtype=float32)
    >>> OP['PG1'].todense()
    matrix([[ 0.],
            [ 0.],
            [ 0.]], dtype=float32)
    
    """
    
    try:
        f=open(pch_file,'r')
    except IOError:
        print "You'd better change the directory of the input file before..."
        raise IOError
    
    lcpt = 0
    
    operators = {}
    
    name = None
    form = 2
    typeIn  = 1
    typeOut = 0
    nRows = None
    nCols = None
    r,c = None,None
    
    line=f.readline()
    
    while line != '':
        
        line = re.match(PCH_PATTERN_Key, line) 
        lcpt+=1
        
        if re.match(r"DMI(?=\*)",line.group(1)): # COLUMN HEADER long format
            entries = line.group(2).split()
            
            name = entries[0]
            try:
                c = int(entries[1])-1
            except ValueError:
                raise "Non matching entry type: COLUMN"
            
            for entry in entries[2:]:
                entry = eval(entry)
                if isinstance(entry,(int,long)):
                    r = entry-1   # explicit change in row number
                elif isinstance(entry,(float)):
                    try:
                        operators[name][r,c] = float(entry)
                    except KeyError:
                        raise "Undefined matrix before line " +  str(lcpt)
                    r += 1      # implicit change in row number
                else:
                    pass
            
        elif re.match(r"DMI(?!\*)",line.group(1)): # HEADER or COLUMN HEADER short format
            entries = line.group(2).split()
            
            if entries[0]!=name:
                
                try:
                    operators[name]=output(operators[name])
                except:
                    pass
                # change operator
#                if form==6: # symmetric matrix
#                    Operators[name] = Operators[name] + Operators[name].transpose() - sp.dia_matrix((Operators[name].diagonal(),0),Operators[name].shape).tolil()
#                    break
#                if form==3:
#                    Operators[name] = sp.dia_matrix((Operators[name].transpose().toarray(),0),(nRows,nRows))
                    
                name = entries[0]
                
            try:
                form = int(entries[2])
            except ValueError:
                raise "Non matching entry type: FORM"
            #typeIn = int(m.group())
            try:
                nRows, nCols = int(entries[5]),int(entries[6])     # get size of the operator
            except ValueError:
                raise "Non matching entry type: Rows and Columns"
            
#            print "Matrix:", name
#            print "nRows, nCols:", isinstance(nRows,(int)), isinstance(nCols,(int))
            
            operators[name] = sp.lil_matrix((nRows,nCols),dtype=FLOAT) # initialize operator
            r,c = None,None # initialize row
            
        elif re.match(r"        ",line.group(1)) or re.match(r"\*       ",line.group(0)): # COLUMN
            entries = line.group(2).split()
            
            for entry in entries:
                entry = eval(entry)
                if isinstance(entry,(int,long)):
                    r = entry-1   # explicit change in row number
                elif isinstance(entry,(float)):
                    try:
                        operators[name][r,c] = entry
                    except KeyError:
                        raise "Undefined matrix before line " +  str(lcpt)
                    r += 1      # implicit change in row number
                else:
                    pass    
        
        
        line=f.readline()
           
    f.close()

    return operators


def writePunch(pch_file,operators,method="overwrite"):
    
    try:
        f=open(pch_file,'w')
    except IOError:
        raise IOError, "Unable to open file", pch_file
    
    LONG = True
    
    for op in operators.keys():
        print "operator:",op
        m,n = operators[op].get_shape()
        # write new head:
        pattern = PCH_HEAD
        pattern["NAME"] = op
        pattern["FORM"] = 2
        pattern["TIN"] = 2
        pattern["TOUT"] = 0
        pattern["M"]= m
        pattern["N"]= n
        hentry = HENTRY
        
        f.write(hentry.format(**pattern))
        
        # write data:

        if LONG:
            entry = ENTRY_LONG
            nentry = 4
        else:
            entry = ENTRY
            nentry = 8
            
            
        index = operators[op].nonzero()
        
        if index[1].size!=0:   
            # get sorted column index
            sortedcol = argsort(index[1])
            
            # number of nonzero item in operator:
            nitem = len(sortedcol)
            
            row=0
            j=0
            while j<nitem:
                
                # start a new column entry:
                # change pattern to column header:
                if LONG:
                    pattern = PCH_HCOLUMN_LONG
                else:
                    pattern = PCH_HCOLUMN
                data=[]
                # set new column head:
                col = index[1][sortedcol[j]]
                data.extend([op,col+1])
                
                # for all nonzero item in column J:
                while 1:
                    
                    
                    if len(data)<nentry:
                        # next row item:
                        if index[0][sortedcol[j]]==row+1:
                            # next row is current row + 1:
                            row+=1
                            data.append(operators[op][row,col])
                            
                        else:
                            # next row is not current row + 1:
                            row = index[0][sortedcol[j]]
                            data.extend([row+1,operators[op][row,col]]) 
                    
                    # ready to write: the number of item in data must be at least equal to
                    # the number of item in the column entry:
                    if len(data)>=nentry:
                        f.write(entry.format(data[0:nentry],**pattern))
                        if len(data)>nentry:
                            data=data[nentry:] # keep unwritten content
                        else:
                            data=[]
                        # change pattern to column pattern:
                        if LONG:
                            pattern = PCH_COLUMN_LONG
                        else:
                            pattern = PCH_COLUMN
                    
                    # get next item:
                    j+=1
                    
                    # Continue loop only if next item exists and
                    # has the same column index.
                    if j >= nitem or col!=index[1][sortedcol[j]]:
                        # going to stop loop over rows...
                        
                        # write residual data of the current column:
                        if len(data)!=0:
                            # fill data with empty string
                            data.extend([" "]*(nentry-len(data)))
                            # write data into file... 
                            f.write(entry.format(data[0:nentry],**pattern)) 
                        
                        break
                        
                
            
    return 0

if __name__ == "__main__":
    
    #import doctest
    #doctest.testmod()
    filename = "../test/beam-ddm.pch"
    
    operators = readPunch(filename,output=sp.lil_matrix)
    
#    UO1 = sp.csr_matrix([[1.0],[0.0],[2.10],[3.2]]) 
#    KOO1 = sp.csr_matrix([[1.0,3.0,0.0,0.0],[0.0,2.0,4.0,0.0],[5.0,0.0,2.0,3.0],[9.0,0.0,0.0,0.0]])
#    
#    
#    
#    operators = {"UO1":UO1.tolil(),"KOO1":KOO1.tolil()}
    
    writePunch("updated-disp.pch",operators)
    
    operators_bis = readPunch("updated-disp.pch",output=sp.lil_matrix)
    
    for op in operators.keys():
        
        A = operators[op] - operators_bis[op]
        print "Checking operator ", op, ": ", A.nonzero()
        
        
    
    